"""
这是任务管理器
"""
import time
from queue import Queue
import csv
import json
import base64
import pickle
import requests
import os
from pprint import pprint
import datetime

import config
from utils import generate_timestamp_tuple, b64encode


class TaskManager:

    def __init__(self) -> None:
        self.task_queue = Queue()
        from client_app import feature_extraction_local, get_bboxes_and_scores_local, concat_result_local, \
            multiclass_nms_local
        self.local_handlers = [feature_extraction_local,
                               get_bboxes_and_scores_local, concat_result_local, multiclass_nms_local]
        self.cloud_handlers = [feature_extraction_cloud,
                               get_bboxes_and_scores_cloud, concat_result_cloud, multiclass_nms_cloud]
        self.max_task_id = 0
        self.current_progress = -1
        self.start_time = None
        self.cpu_time = 0.0
        self.network_speed = 0.0

        # 对cpu和网络进行测速
        self.bench_mark()

        # 可用的文件后缀
        self.valid_suffix = set()
        self.valid_suffix.add("png")
        self.valid_suffix.add("jpg")
        self.valid_suffix.add("jpeg")

        self.csv_path = os.path.join(config.save_path, "result.csv")

        with open(self.csv_path, "w") as f:
            writer = csv.writer(f)
            writer.writerow(['Task ID', 'FileName', 'Action',
                             'Location', 'Function', 'Status', 'Timestamp'])

    def bench_mark(self):
        print("正在测试处理器速度...")
        import math
        # 对cpu进行测速
        st = time.time()
        tmp = 1000005
        for i in range(1, 1000000):
            tmp = math.log2(tmp) * tmp
        et = time.time()
        self.cpu_time = et - st
        print("正在测试网络速度...")
        r = requests.get(config.base_url + 'speedtest/')
        # 计算网速是多少KBytes/s
        print(len(r.content))
        self.network_speed = (1e6 * len(r.content) / r.elapsed.microseconds) / 1024.0
        print("测速完成, 网速为：{} kBytes/s", self.network_speed)
        return self.cpu_time, self.network_speed

    def make_decision(self, do_benchmark=False):
        print("正在作出调度决策")
        if do_benchmark:
            self.bench_mark()

        estimate_local_feature_ext_time = 0.0
        for i, x in enumerate(config.PERFORMANCE_TABLE):
            if self.cpu_time > x[0]:
                estimate_local_feature_ext_time = x[1]
                break
        if estimate_local_feature_ext_time == 0.0:
            estimate_local_feature_ext_time = config.PERFORMANCE_TABLE[0][1]

        estimate_cloud_feature_ext_time = 2048 / self.network_speed + config.SERVER_FEATURE_EXT_TIME

        print("关键步骤 本地执行预计预计耗时：", estimate_local_feature_ext_time)
        print("关键步骤 云端执行预计预计耗时：(网络+计算）", estimate_cloud_feature_ext_time)
        # 云端执行特征提取
        if estimate_cloud_feature_ext_time < estimate_local_feature_ext_time:
            print("决策：[True, False, False, False]")
            return [True, False, False, False]
        else:
            print("决策：[False, False, False, False]")
            # 本地执行特征提取
            return [False, False, False, False]

    def write_csv(self, task_id, path, action, timestamps):
        with open(self.csv_path, "a") as f:
            writer = csv.writer(f)
            for x in timestamps:
                writer.writerow([task_id, path, action, x[0], x[1], x[2], x[3]])

    def add_task(self, file_path):
        self.add_tasks([(tuple(self.make_decision()), file_path)])

    def add_tasks(self, task_list: list):
        for task in task_list:

            if len(task[1].split(".")) > 1 and task[1].split(".")[-1].lower() in self.valid_suffix:
                self.task_queue.put((self.max_task_id, task))
                self.max_task_id += 1
            else:
                print("Ignore: ", task[1])

    def update_progress_bar(self):
        self.current_progress += 1
        print("{}%  ({}/{})".format(100.0 * (1.0 * self.current_progress) /
                                    self.max_task_id, self.current_progress, self.max_task_id))

    def run(self):
        self.start_time = datetime.datetime.now()
        self.update_progress_bar()
        while not self.task_queue.empty():
            t = self.task_queue.get()
            self.__do_one_task(t[0], t[1])
            self.update_progress_bar()

    def __do_one_task(self, task_id, task):
        import client_app

        # 动作标志位
        action, path = task

        # 记录上一步是否在云端执行
        last_cloud = False
        last_data = {
            "result": {
                "img_data": None
            },
            "timestamps": [],
            "action": action
        }

        last_data['timestamps'].append(
            generate_timestamp_tuple("Load img", "begin"))
        print(path)
        _, img_data, img_scale = client_app.load_img(path)
        last_data['result']['img_data'] = b64encode(img_data)
        last_data['result']['img_scale'] = b64encode(img_scale)
        last_data['timestamps'].append(
            generate_timestamp_tuple("Load img", "end"))

        for i, flag in enumerate(action):
            if flag == False:
                # print("i={}, local".format(i))
                last_data = self.local_handlers[i](last_data)
                last_cloud = False
            else:
                # 当前步骤应当在云端执行, 上一步在本地执行
                if last_cloud == False:
                    # print("i={}, cloud".format(i))
                    last_data = self.cloud_handlers[i](last_data)
                    last_cloud = True
        from draw_rec import save_results
        last_data = save_results(last_data, path, draw_thresh=0.5)

        # pprint(last_data['timestamps'])
        self.write_csv(task_id, path, action, last_data['timestamps'])


def feature_extraction_cloud(last_result):
    res = requests.post(url=config.base_url + 'feature_extraction/',
                        headers=config.headers, json=last_result)
    result = res.json()
    return result


def get_bboxes_and_scores_cloud(last_result):
    res = requests.post(url=config.base_url + 'get_bboxes_and_scores/', headers=config.headers,
                        json=last_result)
    result = res.json()
    return result


def concat_result_cloud(last_result):
    res = requests.post(url=config.base_url + 'concat_result/', headers=config.headers,
                        json=last_result)
    result = res.json()
    return result


def multiclass_nms_cloud(last_result):
    res = requests.post(url=config.base_url + 'multiclass_nms/', headers=config.headers,
                        json=last_result)
    result = res.json()
    return result
